#!/usr/bin/env python3
# check_resultsql.py
# usage: 执行后，问题数据写入 tmp_err_sql.log ，可以手工copy到question.jsonl重新运行程序, 运行完毕后，tmp_err_sql.log内容copy到 question.jsonl重新生成result
# python ./src/tools/check_resultsql.py --result "./result/results.jsonl" --question "./data/questions.jsonl"

import argparse
import json
import sqlite3
import os
from typing import Any, Dict, Optional

def parse_args():
    parser = argparse.ArgumentParser(description="检查SQL执行结果，若出错或为空则输出问题和SQL到日志")
    parser.add_argument(
        '--result',
        default='./result/results.jsonl',
        help='结果SQL文件路径，默认 ./data/results.jsonl'
    )
    parser.add_argument(
        '--question',
        default='./data/questions.jsonl',
        help='问题文件路径，默认 ./data/questions.jsonl'
    )
    return parser.parse_args()

def load_jsonl(file_path: str) -> Dict[int, dict]:
    """加载 JSONL 文件，以 id 为 key 构建字典"""
    data = {}
    if not os.path.exists(file_path):
        raise FileNotFoundError(f"文件未找到: {file_path}")
    with open(file_path, 'r', encoding='utf-8') as f:
        for line_num, line in enumerate(f, 1):
            line = line.strip()
            if not line:
                continue
            try:
                item = json.loads(line)
                data[item['id']] = item
            except Exception as e:
                print(f"解析 {file_path} 第 {line_num} 行失败: {e}")
    return data

def execute_sql(cursor: sqlite3.Cursor, sql: str) -> Optional[Any]:
    """执行 SQL，返回结果的第一行第一列，出错返回 None"""
    try:
        cursor.execute(sql)
        row = cursor.fetchone()
        return row[0] if row else None
    except Exception as e:
        print(f"SQL 执行错误: {e}")
        return None

def main():
    args = parse_args()

    # 加载数据
    try:
        result_data = load_jsonl(args.result)
        question_data = load_jsonl(args.question)
    except Exception as e:
        print(f"加载文件失败: {e}")
        return

    db_path = './data/data.db'
    if not os.path.exists(db_path):
        print(f"数据库文件不存在: {db_path}")
        return

    # 连接数据库
    conn = sqlite3.connect(db_path)
    cursor = conn.cursor()

    # 打开日志文件
    log_file = 'tmp_err_sql.log'
    with open(log_file, 'w', encoding='utf-8') as log_f:
        error_count = 0

        # 按 id 顺序处理 result 文件中的每条 SQL
        for idx in sorted(result_data.keys()):
            item = result_data[idx]
            sql = item.get('sql', '').strip()
            qid = item['id']

            # 查找对应的问题
            question_item = question_data.get(qid)
            if not question_item:
                print(f"警告: result 中 id={qid} 在 question 文件中未找到")
                continue

            # 执行 SQL
            result = execute_sql(cursor, sql)

            # 判断是否出错或结果为空
            if result is None:
                # 构造输出信息
                output = {
                    "id": qid,
                    "question": question_item['question'],
                    "level": question_item['level'],
                    "sql": sql
                }
                output_str = json.dumps(output, ensure_ascii=False)
                print("\n[ERROR] SQL 执行失败或结果为空:")
                print(output_str)

                # 写入日志
                log_f.write(output_str + '\n')
                error_count += 1

        conn.close()

        # 最终统计
        print(f"\n检查完成。共发现 {error_count} 条错误/空结果 SQL，已写入 {log_file}")

if __name__ == '__main__':
    main()